;;; Longest Collatz sequence

;;; Problem 14

;;; The following iterative sequence is defined for the set
;;; of positive integers:

;;; n -> n/2 (n is even)
;;; n -> 3n + 1 (n is odd)

;;; Using the rule above and starting with 13, we generate
;;; the following sequence:

;;; 13 -> 40 -> 20 -> 10 -> 5 -> 16 -> 8 -> 4 -> 2 -> 1

;;; It can be seen that this sequence (starting at 13 and
;;; finishing at 1) contains 10 terms. Although it has not
;;; been proved yet (Collatz Problem), it is thought that
;;; all starting numbers finish at 1.

;;; Which starting number, under one million, produces the
;;; longest chain?

;;; NOTE: Once the chain starts the terms are allowed to go
;;; above one million.

(import
 (except (rnrs base) let-values map)
 (only (guile)
       lambda* λ
       ;; printing
       display
       simple-format)
 (ice-9 futures)
 (srfi srfi-69)  ; hash tables
 (srfi srfi-1)  ; reduce
 (lib math)
 (lib segment))


(define collatz-step
  (λ (num)
    (cond
     [(even? num) (/ num 2)]
     [else (+ (* 3 num) 1)])))


(define collatz-sequence-length
  (λ (seq-start-num)
    ;; (display (simple-format #f "Start of sequence: ~a\n" seq-start-num))
    (let ([seen-numbers (make-hash-table =)])
      (let loop ([sequence-index 1] [num-in-seq seq-start-num])
        ;; (display (simple-format #f "Number in sequence: ~a\n" num-in-seq))
        (cond
         ;; If the number has already been seen, stop and
         ;; return the sequence length.
         [(hash-table-ref/default seen-numbers num-in-seq #f)
          sequence-index]
         [else
          (hash-table-set! seen-numbers num-in-seq #t)
          (loop (+ sequence-index 1)
                (collatz-step num-in-seq))])))))


(define find-longest-collatz-sequence
  (λ (start limit)
    ;; Calculate the maximum of sequence lengths for all
    ;; numbers from start to limit.
    (let iter-sequence-start ([seq-start start]
                              [longest-seq-len 0]
                              [number-with-longest-seq 0])
      (cond
       [(<= seq-start limit)
        (let ([seq-len (collatz-sequence-length seq-start)])
          (cond
           [(> seq-len longest-seq-len)
            ;; (display (simple-format #f "found new longest with length: ~a\n" seq-len))
            (iter-sequence-start (+ seq-start 1)
                                 seq-len
                                 seq-start)]
           [else
            (iter-sequence-start (+ seq-start 1)
                                 longest-seq-len
                                 number-with-longest-seq)]))]
       [else
        (display
         (simple-format
          #f "number with longest sequence in segment ~a-~a: ~a (with length ~a)\n"
          start limit number-with-longest-seq longest-seq-len))
        (cons number-with-longest-seq longest-seq-len)]))))


(define run-in-parallel
  (λ (segments map-proc reduce-proc reduce-init)
    "Use futures to run a procedure in parallel, if multiple
cores are available. Take a list of SEGMENTS as input, which
are ranges of values to work on. MAP-PROC is applied to the
SEGMENTS using map. When the MAP-PROC calls for all segments
finished and returned values, the REDUCE-PROC is applied to
the map result using reduce and the REDUCE-INIT argument."
    (let ([futures
           (map (λ (seg)
                  (make-future
                   ;; Need to wrap in a thunk, to not
                   ;; immediately start evaluating.
                   (λ () (map-proc seg))))
                segments)])
      (let ([segment-results (map touch futures)])
        (reduce reduce-proc
                reduce-init
                segment-results)))))


(let* ([start 1]
       [end (expt 10 6)]
       [num-cores 64]
       [segments (segment start end num-cores)])
  (display
   (simple-format
    #f "longest sequence length (number . length): ~a\n"
    (run-in-parallel segments
                     (λ (seg)
                       (find-longest-collatz-sequence (segment-start seg)
                                                      (segment-end seg)))
                     (λ (acc current)
                       (if (> (cdr acc) (cdr current))
                           acc
                           current))
                     -inf.0))))
